import copy
import torch
import gc
import math
import torch.nn as nn
from fast_pytorch_kmeans import KMeans

def weight_fed(global_protos, local_protos, opt):
    protos_bank = []
    #将所有非NaN的本地记忆拼成一个protos_bank，待会用于聚类，local_protos是一个字典，它的键是客户端ID
    for k, v in local_protos.items():
        if bool(torch.isnan(v[0]).any()):
            continue
        else:
            protos_bank.append(v[0])
    protos_bank = torch.cat(protos_bank, dim=0).contiguous()#按第0维拼接成torch张量
    protos_bank = protos_bank.cuda()#移动到gpu上
    protos_bank = torch.cat((global_protos, protos_bank), dim=0) #将本地与全局记忆拼在一起

    kmeans = KMeans(n_clusters=global_protos.shape[0], mode='cosine') # 聚类中心数等于全局记忆的个数使用余弦相似度作为距离度量
    # 对 protos_bank 进行K-means聚类，返回 cluster_r，其中每个元素代表对应原型的聚类结果。
    cluster_r = kmeans.fit_predict(protos_bank)
    # 包含全局记忆的聚类标签。
    cluster_global = cluster_r[0:global_protos.shape[0]]
    # 包含本地记忆的聚类标签。
    cluster_local = cluster_r[global_protos.shape[0]:len(cluster_r)]

    #创建一个稀疏单位矩阵 ones（形状为全局原型数目×全局原型数目）并放入GPU。
    ones = torch.sparse.torch.eye(global_protos.shape[0]).cuda()
    #根据 cluster_global 中的聚类标签选择 ones 中的对应行，生成 mm_global（矩阵的子集）。
    mm_global = ones.index_select(0, cluster_global)

    #存储各个本地模型的权重。
    modal_w = torch.zeros(1, len(local_protos))
    softMax = nn.Softmax(dim=1)

    #遍历 local_protos 字典，对每个客户端的本地原型计算权重。
    for k, v in local_protos.items():
        cluster_label = cluster_local[k * global_protos.shape[0]: (k+1) * global_protos.shape[0]].cuda()
        mm_local = ones.index_select(0, cluster_label)

        #度量全局和本地原型之间的相似性theta
        theta = 1.0 / 2 * torch.matmul(global_protos, v[0].cuda().t())
        #计算本地记忆和全局记忆的相似度Sim
        S = torch.matmul(mm_global, mm_local.t())
        if opt.use_gpu:
            S = S.cuda()

        #根据 theta 和 S 计算客户端 k 的权重，并将结果存储在 modal_w 的第 k 列中。
        modal_w[0][k] = torch.log(-torch.sum(S * theta - torch.log(1.0 + torch.exp(theta))))

    modal_w = softMax(modal_w)
    return modal_w

def weighted_agg_fed(w, modal_w):
    w_avg = copy.deepcopy(w[0])
    for k in w_avg.keys():
        w_avg[k] = w_avg[k] * modal_w[0][0]
        for i in range(1, len(w)):
            w_avg[k] += w[i][k] * modal_w[0][i]

    return w_avg



